"""
U-Net模型定义

这个文件包含了U-Net模型的架构定义，用于医学图像分割任务。
U-Net架构包括编码器(下采样路径)和解码器(上采样路径)，以及它们之间的跳跃连接。
"""

import tensorflow as tf
from keras.layers import Input, Conv2D, BatchNormalization, Activation, MaxPool2D
from keras.layers import UpSampling2D, Concatenate
from keras.models import Model

def conv_block(x, num_filters):
    """
    卷积块: 两个卷积层，每层后面跟着批归一化和ReLU激活函数
    
    参数:
        x: 输入张量
        num_filters: 卷积层的滤波器数量
    返回:
        x: 处理后的张量
    """
    x = Conv2D(num_filters, (3, 3), padding="same")(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    
    x = Conv2D(num_filters, (3, 3), padding="same")(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    
    return x

def build_unet_model(input_shape):
    """
    构建U-Net模型
    
    参数:
        input_shape: 输入图像的形状 (高度, 宽度, 通道数)
    返回:
        model: 编译好的Keras模型
    """
    # 定义每个编码器层的滤波器数量
    num_filters = [64, 128, 256, 512]
    
    # 输入层
    inputs = Input(input_shape)
    
    # 存储跳跃连接的列表
    skip_connections = []
    
    # 编码器路径(下采样)
    x = inputs
    for filters in num_filters:
        x = conv_block(x, filters)  # 应用卷积块
        skip_connections.append(x)  # 存储用于跳跃连接的特征图
        x = MaxPool2D((2, 2))(x)   # 下采样
    
    # 瓶颈(网络最深处)
    x = conv_block(x, 1024)
    
    # 准备解码器路径
    num_filters.reverse()  # 反转滤波器列表用于解码器
    skip_connections.reverse()  # 反转跳跃连接列表
    
    # 解码器路径(上采样)
    for i, filters in enumerate(num_filters):
        x = UpSampling2D((2, 2))(x)  # 上采样
        x = Concatenate()([x, skip_connections[i]])  # 连接跳跃连接
        x = conv_block(x, filters)  # 应用卷积块
    
    # 输出层 - 使用sigmoid激活函数进行二元分割
    outputs = Conv2D(1, (1, 1), padding="same", activation="sigmoid")(x)
    
    # 创建并返回模型
    model = Model(inputs, outputs)
    return model

if __name__ == "__main__":
    # 测试模型构建
    input_shape = (128, 128, 3)  # 高度x宽度x通道数
    model = build_unet_model(input_shape)
    model.summary()  # 打印模型架构
